#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2024/12/3 15:41
@Author  : lxc
@File    : sqlUtil.py
@Desc    :

"""
import sqlite3
import logging

class ItemToSQLite:
    """
    将数据存储到SQLite数据库中
    """

    def __init__(self, db_path, tb_name):
        self.db_path = db_path  # SQLite数据库文件路径
        self.tb_name = tb_name  # 表名

    def connectSQLite(self):
        """
        连接到SQLite数据库
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            return conn, cursor
        except Exception as err:
            logging.error("数据库连接失败, %s" % err)
            return None, None

    def createSQLite(self, meta, table_name):
        """
        根据字段名创建SQLite表
        """
        conn, cursor = self.connectSQLite()
        keys = meta.keys()
        sql = f"CREATE TABLE IF NOT EXISTS `{table_name}` (id INTEGER PRIMARY KEY AUTOINCREMENT, "
        sql += ", ".join([f"`{key}` TEXT" for key in keys]) + ");"

        try:
            cursor.execute(sql)
            conn.commit()
            logging.info(f"创建表 {table_name} 成功")
        except Exception as err:
            conn.rollback()
            logging.error(f"创建表 {table_name} 失败, 已回滚! {err}")
        finally:
            conn.close()

    def insertSQLite(self, meta, tb_name):
        """
        插入数据到SQLite表
        """
        conn, cursor = self.connectSQLite()
        keys = list(meta.keys())
        values = [meta[key] if meta[key] else "" for key in keys]
        # sql = f"INSERT INTO `{tb_name}` ({', '.join(keys)}) VALUES ({', '.join(['?'] * len(values))})"
        sql = f"INSERT INTO `{tb_name}` ({', '.join([f'`{key}`' for key in keys])}) VALUES ({', '.join(['?'] * len(values))})"

        try:
            cursor.execute(sql, values)
            conn.commit()
            logging.info(f"插入数据 {meta} 成功")
        except Exception as err:
            conn.rollback()
            logging.error(f"插入数据 {meta} 失败, 已回滚! {err}")
            return str(err)
        finally:
            conn.close()

    def add_field(self, meta, tb_name):
        """
        如果表中缺少字段，添加字段
        """
        conn, cursor = self.connectSQLite()
        current_columns = self.query_columns(tb_name)
        fields_to_add = [field for field in meta.keys() if field not in current_columns]

        for field in fields_to_add:
            sql = f"ALTER TABLE `{tb_name}` ADD COLUMN `{field}` TEXT"
            try:
                cursor.execute(sql)
                conn.commit()
                logging.info(f"表 {tb_name} 添加字段 {field} 成功")
            except Exception as err:
                conn.rollback()
                logging.error(f"表 {tb_name} 添加字段 {field} 失败, 已回滚! {err}")
        conn.close()

    def query_columns(self, tb_name):
        """
        查询表中已有的字段名
        """
        conn, cursor = self.connectSQLite()
        sql = f"PRAGMA table_info({tb_name})"
        columns = []
        try:
            cursor.execute(sql)
            columns = [column[1] for column in cursor.fetchall()]
            logging.info(f"{tb_name} 表字段: {columns}")
        except Exception as err:
            logging.error(f"查询表 {tb_name} 字段失败: {err}")
        finally:
            conn.close()
        return columns

    def process_item(self, item):
        """
        处理数据项：插入数据，表不存在时创建，字段不一致时添加
        """
        err = self.insertSQLite(item, self.tb_name)

        # 如果插入失败并提示表不存在，则创建表并重新插入数据
        if "no such table" in str(err):
            self.createSQLite(item, self.tb_name)
            self.insertSQLite(item, self.tb_name)
        # 如果插入失败并提示字段不一致，则添加字段并重新插入数据
        elif "no column" in str(err):
            self.add_field(item, self.tb_name)
            self.insertSQLite(item, self.tb_name)

if __name__ == "__main__":
    # 创建对象并处理数据
    store = ItemToSQLite(db_path="db.db", tb_name="industry_info")
    data = {
        "title": "行业新闻",
        "static_url": "",
        "abstract": "行业新闻摘要",
        "author": "张三"
    }
    store.process_item(item=data)
